
/*****************************************************************************
*
* MODULE:       Grass PDE Numerical Library
* AUTHOR(S):    Soeren Gebbert, Berlin (GER) Dec 2006
* 		soerengebbert <at> gmx <dot> de
*               
* PURPOSE:      groundwater flow in porous media 
* 		part of the gpde library
*
* COPYRIGHT:    (C) 2000 by the GRASS Development Team
*
*               This program is free software under the GNU General Public
*               License (>=v2). Read the file COPYING that comes with GRASS
*               for details.
*
*****************************************************************************/

#include "grass/N_gwflow.h"

/* *************************************************************** */
/* ***************** N_gwflow_data3d ***************************** */
/* *************************************************************** */
/*!
 * \brief Alllocate memory for the groundwater calculation data structure in 3 dimensions
 *
 * The groundwater calculation data structure will be allocated including
 * all appendant 3d and 2d arrays. The offset for the 3d arrays is one
 * to establish homogeneous Neumann boundary conditions at the calculation area border.
 * This data structure is used to create a linear equation system based on the computation of
 * groundwater flow in porous media with the finite volume method.
 *
 * \param cols   int
 * \param rows   int
 * \param depths int
 * \return N_gwflow_data3d *
 * */
N_gwflow_data3d *N_alloc_gwflow_data3d(int cols, int rows, int depths,
				       int river, int drain)
{
    N_gwflow_data3d *data;

    data = (N_gwflow_data3d *) G_calloc(1, sizeof(N_gwflow_data3d));

    data->phead = N_alloc_array_3d(cols, rows, depths, 1, DCELL_TYPE);
    data->phead_start = N_alloc_array_3d(cols, rows, depths, 1, DCELL_TYPE);
    data->status = N_alloc_array_3d(cols, rows, depths, 1, DCELL_TYPE);
    data->hc_x = N_alloc_array_3d(cols, rows, depths, 1, DCELL_TYPE);
    data->hc_y = N_alloc_array_3d(cols, rows, depths, 1, DCELL_TYPE);
    data->hc_z = N_alloc_array_3d(cols, rows, depths, 1, DCELL_TYPE);
    data->q = N_alloc_array_3d(cols, rows, depths, 1, DCELL_TYPE);
    data->s = N_alloc_array_3d(cols, rows, depths, 1, DCELL_TYPE);
    data->nf = N_alloc_array_3d(cols, rows, depths, 1, DCELL_TYPE);
    data->r = N_alloc_array_2d(cols, rows, 1, DCELL_TYPE);

    if (river) {
	data->river_head = N_alloc_array_3d(cols, rows, depths, 1, DCELL_TYPE);
	data->river_leak = N_alloc_array_3d(cols, rows, depths, 1, DCELL_TYPE);
	data->river_bed = N_alloc_array_3d(cols, rows, depths, 1, DCELL_TYPE);
    }
    else {
	data->river_head = NULL;
	data->river_leak = NULL;
	data->river_bed = NULL;
    }

    if (drain) {
	data->drain_leak = N_alloc_array_3d(cols, rows, depths, 1, DCELL_TYPE);
	data->drain_bed = N_alloc_array_3d(cols, rows, depths, 1, DCELL_TYPE);
    }
    else {
	data->drain_leak = NULL;
	data->drain_bed = NULL;
    }

    return data;
}

/* *************************************************************** */
/* ********************* N_free_gwflow_data3d ******************** */
/* *************************************************************** */
/*!
 * \brief Release the memory of the groundwater flow data structure in three dimensions
 *
 * \param data N_gwflow_data3d *
 * \return void *
 * */

void N_free_gwflow_data3d(N_gwflow_data3d * data)
{
    if (data->phead)
	N_free_array_3d(data->phead);
    if (data->phead_start)
	N_free_array_3d(data->phead_start);
    if (data->status)
	N_free_array_3d(data->status);
    if (data->hc_x)
	N_free_array_3d(data->hc_x);
    if (data->hc_y)
	N_free_array_3d(data->hc_y);
    if (data->hc_z)
	N_free_array_3d(data->hc_z);
    if (data->q)
	N_free_array_3d(data->q);
    if (data->s)
	N_free_array_3d(data->s);
    if (data->nf)
	N_free_array_3d(data->nf);
    if (data->r)
	N_free_array_2d(data->r);
    if (data->river_head)
	N_free_array_3d(data->river_head);
    if (data->river_leak)
	N_free_array_3d(data->river_leak);
    if (data->river_bed)
	N_free_array_3d(data->river_bed);
    if (data->drain_leak)
	N_free_array_3d(data->drain_leak);
    if (data->drain_bed)
	N_free_array_3d(data->drain_bed);

    G_free(data);

    data = NULL;

    return;
}

/* *************************************************************** */
/* ******************** N_alloc_gwflow_data2d ******************** */
/* *************************************************************** */
/*!
 * \brief Alllocate memory for the groundwater calculation data structure in 2 dimensions
 * 
 * The groundwater calculation data structure will be allocated including
 * all appendant 2d arrays. The offset for the 3d arrays is one
 * to establish homogeneous Neumann boundary conditions at the calculation area border.
 * This data structure is used to create a linear equation system based on the computation of
 * groundwater flow in porous media with the finite volume method.
 *
 * \param cols int
 * \param rows int
 * \return N_gwflow_data2d *
 * */
N_gwflow_data2d *N_alloc_gwflow_data2d(int cols, int rows, int river, int drain)
{
    N_gwflow_data2d *data;

    data = (N_gwflow_data2d *) G_calloc(1, sizeof(N_gwflow_data2d));

    data->phead = N_alloc_array_2d(cols, rows, 1, DCELL_TYPE);
    data->phead_start = N_alloc_array_2d(cols, rows, 1, DCELL_TYPE);
    data->status = N_alloc_array_2d(cols, rows, 1, DCELL_TYPE);
    data->hc_x = N_alloc_array_2d(cols, rows, 1, DCELL_TYPE);
    data->hc_y = N_alloc_array_2d(cols, rows, 1, DCELL_TYPE);
    data->q = N_alloc_array_2d(cols, rows, 1, DCELL_TYPE);
    data->s = N_alloc_array_2d(cols, rows, 1, DCELL_TYPE);
    data->nf = N_alloc_array_2d(cols, rows, 1, DCELL_TYPE);
    data->r = N_alloc_array_2d(cols, rows, 1, DCELL_TYPE);
    data->top = N_alloc_array_2d(cols, rows, 1, DCELL_TYPE);
    data->bottom = N_alloc_array_2d(cols, rows, 1, DCELL_TYPE);

    if (river) {
	data->river_head = N_alloc_array_2d(cols, rows, 1, DCELL_TYPE);
	data->river_leak = N_alloc_array_2d(cols, rows, 1, DCELL_TYPE);
	data->river_bed = N_alloc_array_2d(cols, rows, 1, DCELL_TYPE);
    }
    else {
	data->river_head = NULL;
	data->river_leak = NULL;
	data->river_bed = NULL;
    }

    if (drain) {
	data->drain_leak = N_alloc_array_2d(cols, rows, 1, DCELL_TYPE);
	data->drain_bed = N_alloc_array_2d(cols, rows, 1, DCELL_TYPE);
    }
    else {
	data->drain_leak = NULL;
	data->drain_bed = NULL;
    }


    return data;
}

/* *************************************************************** */
/* ****************** N_free_gwflow_data2d *********************** */
/* *************************************************************** */
/*!
 * \brief Release the memory of the groundwater flow data structure in two dimensions
 *
 * \param data N_gwflow_data2d *
 * \return void
 * */
void N_free_gwflow_data2d(N_gwflow_data2d * data)
{
    if (data->phead)
	N_free_array_2d(data->phead);
    if (data->phead_start)
	N_free_array_2d(data->phead_start);
    if (data->status)
	N_free_array_2d(data->status);
    if (data->hc_x)
	N_free_array_2d(data->hc_x);
    if (data->hc_y)
	N_free_array_2d(data->hc_y);
    if (data->q)
	N_free_array_2d(data->q);
    if (data->s)
	N_free_array_2d(data->s);
    if (data->nf)
	N_free_array_2d(data->nf);
    if (data->r)
	N_free_array_2d(data->r);
    if (data->top)
	N_free_array_2d(data->top);
    if (data->bottom)
	N_free_array_2d(data->bottom);
    if (data->river_head)
	N_free_array_2d(data->river_head);
    if (data->river_leak)
	N_free_array_2d(data->river_leak);
    if (data->river_bed)
	N_free_array_2d(data->river_bed);
    if (data->drain_leak)
	N_free_array_2d(data->drain_leak);
    if (data->drain_bed)
	N_free_array_2d(data->drain_bed);

    G_free(data);

    data = NULL;;

    return;
}

/* *************************************************************** */
/* ***************** N_callback_gwflow_3d ************************ */
/* *************************************************************** */
/*!
 * \brief This callback function creates the mass balance of a 7 point star
 *
 * The mass balance is based on the common groundwater flow equation:
 *
 * \f[Ss \frac{\partial h}{\partial t} = \nabla {\bf K} \nabla h + q \f]
 *
 * This equation is discretizised with the finite volume method in three dimensions.
 *
 *
 * \param gwdata N_gwflow_data3d *
 * \param geom N_geom_data *
 * \param col   int
 * \param row   int
 * \param depth int
 * \return N_data_star *
 *
 * */
N_data_star *N_callback_gwflow_3d(void *gwdata, N_geom_data * geom, int col,
				  int row, int depth)
{
    double hc_e = 0, hc_w = 0, hc_n = 0, hc_s = 0, hc_t = 0, hc_b = 0;
    double dx, dy, dz, Ax, Ay, Az;
    double hc_x, hc_y, hc_z;
    double hc_xw, hc_yn, hc_zt;
    double hc_xe, hc_ys, hc_zb;
    double hc_start;
    double Ss, r, nf, q;
    double C, W, E, N, S, T, B, V;
    N_data_star *mat_pos;
    N_gwflow_data3d *data;

    /*cast the void pointer to the right data structure */
    data = (N_gwflow_data3d *) gwdata;

    dx = geom->dx;
    dy = geom->dy;
    dz = geom->dz;
    Az = N_get_geom_data_area_of_cell(geom, row);
    Ay = geom->dx * geom->dz;
    Ax = geom->dz * geom->dy;

    /*read the data from the arrays */
    hc_start = N_get_array_3d_d_value(data->phead_start, col, row, depth);

    hc_x = N_get_array_3d_d_value(data->hc_x, col, row, depth);
    hc_y = N_get_array_3d_d_value(data->hc_y, col, row, depth);
    hc_z = N_get_array_3d_d_value(data->hc_z, col, row, depth);

    hc_xw = N_get_array_3d_d_value(data->hc_x, col - 1, row, depth);
    hc_xe = N_get_array_3d_d_value(data->hc_x, col + 1, row, depth);
    hc_yn = N_get_array_3d_d_value(data->hc_y, col, row - 1, depth);
    hc_ys = N_get_array_3d_d_value(data->hc_y, col, row + 1, depth);
    hc_zt = N_get_array_3d_d_value(data->hc_z, col, row, depth + 1);
    hc_zb = N_get_array_3d_d_value(data->hc_z, col, row, depth - 1);

    hc_w = N_calc_harmonic_mean(hc_xw, hc_x);
    hc_e = N_calc_harmonic_mean(hc_xe, hc_x);
    hc_n = N_calc_harmonic_mean(hc_yn, hc_y);
    hc_s = N_calc_harmonic_mean(hc_ys, hc_y);
    hc_t = N_calc_harmonic_mean(hc_zt, hc_z);
    hc_b = N_calc_harmonic_mean(hc_zb, hc_z);

    /*inner sources */
    q = N_get_array_3d_d_value(data->q, col, row, depth);
    /*specific yield */
    Ss = N_get_array_3d_d_value(data->s, col, row, depth);
    /*porosity */
    nf = N_get_array_3d_d_value(data->nf, col, row, depth);

    /*mass balance center cell to western cell */
    W = -1 * Ax * hc_w / dx;
    /*mass balance center cell to eastern cell */
    E = -1 * Ax * hc_e / dx;
    /*mass balance center cell to northern cell */
    N = -1 * Ay * hc_n / dy;
    /*mass balance center cell to southern cell */
    S = -1 * Ay * hc_s / dy;
    /*mass balance center cell to top cell */
    T = -1 * Az * hc_t / dz;
    /*mass balance center cell to bottom cell */
    B = -1 * Az * hc_b / dz;

    /*specific yield */
    Ss = Az * dz * Ss;

    /*the diagonal entry of the matrix */
    C = -1 * (W + E + N + S + T + B - Ss / data->dt * Az);

    /*the entry in the right side b of Ax = b */
    V = (q + hc_start * Ss / data->dt * Az);

    /*only the top cells will have recharge */
    if (depth == geom->depths - 2) {
	r = N_get_array_2d_d_value(data->r, col, row);
	V += r * Az;
    }

    G_debug(5, "N_callback_gwflow_3d: called [%i][%i][%i]", depth, col, row);

    /*create the 7 point star entries */
    mat_pos = N_create_7star(C, W, E, N, S, T, B, V);

    return mat_pos;
}

/* *************************************************************** */
/* ****************** N_callback_gwflow_2d *********************** */
/* *************************************************************** */
/*!
 * \brief This callback function creates the mass balance of a 5 point star
 *
 * The mass balance is based on the common groundwater flow equation:
 *
 * \f[Ss \frac{\partial h}{\partial t} = \nabla {\bf K} \nabla h + q \f]
 *
 * This equation is discretizised with the finite volume method in two dimensions.
 *
 * \param gwdata N_gwflow_data2d *
 * \param geom N_geom_data *
 * \param col int
 * \param row int
 * \return N_data_star *
 *
 * */
N_data_star *N_callback_gwflow_2d(void *gwdata, N_geom_data * geom, int col,
				  int row)
{
    double T_e = 0, T_w = 0, T_n = 0, T_s = 0;
    double z_e = 0, z_w = 0, z_n = 0, z_s = 0;
    double dx, dy, Az;
    double hc_x, hc_y;
    double z, top;
    double hc_xw, hc_yn;
    double z_xw, z_yn;
    double hc_xe, hc_ys;
    double z_xe, z_ys;
    double hc, hc_start;
    double Ss, r, nf, q;
    double C, W, E, N, S, V;
    N_gwflow_data2d *data;
    N_data_star *mat_pos;
    double river_vect = 0;	/*entry in vector */
    double river_mat = 0;	/*entry in matrix */
    double drain_vect = 0;	/*entry in vector */
    double drain_mat = 0;	/*entry in matrix */

    /*cast the void pointer to the right data structure */
    data = (N_gwflow_data2d *) gwdata;

    dx = geom->dx;
    dy = geom->dy;
    Az = N_get_geom_data_area_of_cell(geom, row);

    /*read the data from the arrays */
    hc_start = N_get_array_2d_d_value(data->phead_start, col, row);
    hc = N_get_array_2d_d_value(data->phead, col, row);
    top = N_get_array_2d_d_value(data->top, col, row);


    if (hc > top) {		/*If the aquifer is confined */
	z = N_get_array_2d_d_value(data->top, col,
				   row) -
	    N_get_array_2d_d_value(data->bottom, col, row);
	z_xw =
	    N_get_array_2d_d_value(data->top, col - 1,
				   row) -
	    N_get_array_2d_d_value(data->bottom, col - 1, row);
	z_xe =
	    N_get_array_2d_d_value(data->top, col + 1,
				   row) -
	    N_get_array_2d_d_value(data->bottom, col + 1, row);
	z_yn =
	    N_get_array_2d_d_value(data->top, col,
				   row - 1) -
	    N_get_array_2d_d_value(data->bottom, col, row - 1);
	z_ys =
	    N_get_array_2d_d_value(data->top, col,
				   row + 1) -
	    N_get_array_2d_d_value(data->bottom, col, row + 1);
    }
    else {			/* the aquifer is unconfined */

	/* If the aquifer is unconfied use an explicite scheme to solve
	 * the nonlinear equation. We use the phead from the first iteration */
	z = N_get_array_2d_d_value(data->phead, col, row) -
	    N_get_array_2d_d_value(data->bottom, col, row);
	z_xw = N_get_array_2d_d_value(data->phead, col - 1, row) -
	    N_get_array_2d_d_value(data->bottom, col - 1, row);
	z_xe = N_get_array_2d_d_value(data->phead, col + 1, row) -
	    N_get_array_2d_d_value(data->bottom, col + 1, row);
	z_yn = N_get_array_2d_d_value(data->phead, col, row - 1) -
	    N_get_array_2d_d_value(data->bottom, col, row - 1);
	z_ys = N_get_array_2d_d_value(data->phead, col, row + 1) -
	    N_get_array_2d_d_value(data->bottom, col, row + 1);
    }

    /*geometrical mean of cell height */
    if (z_w > 0 || z_w < 0 || z_w == 0)
	z_w = N_calc_arith_mean(z_xw, z);
    else
	z_w = z;
    if (z_e > 0 || z_e < 0 || z_e == 0)
	z_e = N_calc_arith_mean(z_xe, z);
    else
	z_e = z;
    if (z_n > 0 || z_n < 0 || z_n == 0)
	z_n = N_calc_arith_mean(z_yn, z);
    else
	z_n = z;
    if (z_s > 0 || z_s < 0 || z_s == 0)
	z_s = N_calc_arith_mean(z_ys, z);
    else
	z_s = z;

    /* Inner sources */
    q = N_get_array_2d_d_value(data->q, col, row);
    nf = N_get_array_2d_d_value(data->nf, col, row);

    /* specific yield  of current cell face */
    Ss = N_get_array_2d_d_value(data->s, col, row) * Az;
    /* recharge */
    r = N_get_array_2d_d_value(data->r, col, row) * Az;

    /*get the surrounding permeabilities */
    hc_x = N_get_array_2d_d_value(data->hc_x, col, row);
    hc_y = N_get_array_2d_d_value(data->hc_y, col, row);
    hc_xw = N_get_array_2d_d_value(data->hc_x, col - 1, row);
    hc_xe = N_get_array_2d_d_value(data->hc_x, col + 1, row);
    hc_yn = N_get_array_2d_d_value(data->hc_y, col, row - 1);
    hc_ys = N_get_array_2d_d_value(data->hc_y, col, row + 1);

    /* calculate the transmissivities */
    T_w = N_calc_harmonic_mean(hc_xw, hc_x) * z_w;
    T_e = N_calc_harmonic_mean(hc_xe, hc_x) * z_e;
    T_n = N_calc_harmonic_mean(hc_yn, hc_y) * z_n;
    T_s = N_calc_harmonic_mean(hc_ys, hc_y) * z_s;

    /*compute the river leakage, this is an explicit method */
    if (data->river_leak &&
	(N_get_array_2d_d_value(data->river_leak, col, row) != 0)) {
	if (hc > N_get_array_2d_d_value(data->river_bed, col, row)) {
	    river_vect = N_get_array_2d_d_value(data->river_head, col, row) *
		N_get_array_2d_d_value(data->river_leak, col, row);
	    river_mat = N_get_array_2d_d_value(data->river_leak, col, row);
	}
	else if (hc < N_get_array_2d_d_value(data->river_bed, col, row)) {
	    river_vect = (N_get_array_2d_d_value(data->river_head, col, row) -
			  N_get_array_2d_d_value(data->river_bed, col, row))
		* N_get_array_2d_d_value(data->river_leak, col, row);
	    river_mat = 0;
	}
    }

    /*compute the drainage, this is an explicit method */
    if (data->drain_leak &&
	(N_get_array_2d_d_value(data->drain_leak, col, row) != 0)) {
	if (hc > N_get_array_2d_d_value(data->drain_bed, col, row)) {
	    drain_vect = N_get_array_2d_d_value(data->drain_bed, col, row) *
		N_get_array_2d_d_value(data->drain_leak, col, row);
	    drain_mat = N_get_array_2d_d_value(data->drain_leak, col, row);
	}
	else if (hc <= N_get_array_2d_d_value(data->drain_bed, col, row)) {
	    drain_vect = 0;
	    drain_mat = 0;
	}
    }

    /*mass balance center cell to western cell */
    W = -1 * T_w * dy / dx;
    /*mass balance center cell to eastern cell */
    E = -1 * T_e * dy / dx;
    /*mass balance center cell to northern cell */
    N = -1 * T_n * dx / dy;
    /*mass balance center cell to southern cell */
    S = -1 * T_s * dx / dy;

    /*the diagonal entry of the matrix */
    C = -1 * (W + E + N + S - Ss / data->dt - river_mat * Az - drain_mat * Az);

    /*the entry in the right side b of Ax = b */
    V = (q + hc_start * Ss / data->dt) + r + river_vect * Az + drain_vect * Az;

    G_debug(5, "N_callback_gwflow_2d: called [%i][%i]", row, col);

    /*create the 5 point star entries */
    mat_pos = N_create_5star(C, W, E, N, S, V);

    return mat_pos;
}
